package com.roy.shardingDemo.algorithm;

import com.google.common.collect.Range;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.sharding.api.sharding.complex.ComplexKeysShardingValue;

import java.util.*;
import java.util.stream.Collectors;

public class OmsOrderShardingAlgorithm implements ComplexKeysShardingAlgorithm<Long> {

    private static final String COLUMN_ORDER_SHARDING_KEY = "id";
    private static final String COLUMN_CUSTOMER_SHARDING_KEY = "member_id";

    @Override
    public Collection<String> doSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Long> shardingValue) {
        // 处理精确查询和IN条件
        if (!shardingValue.getColumnNameAndShardingValuesMap().isEmpty()) {
            return handlePointSharding(availableTargetNames, shardingValue);
        }
        // 处理范围查询
        else if (!shardingValue.getColumnNameAndRangeValuesMap().isEmpty()) {
            return handleRangeSharding(availableTargetNames, shardingValue);
        }
        return availableTargetNames;
    }

    private Collection<String> handlePointSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Long> shardingValue) {
        Set<String> result = new LinkedHashSet<>();
        Map<String, Collection<Long>> columnShardingValues = shardingValue.getColumnNameAndShardingValuesMap();

        List<String> orderIds = convertToStringList(columnShardingValues.getOrDefault(COLUMN_ORDER_SHARDING_KEY, Collections.emptyList()));
        List<String> customerIds = convertToStringList(columnShardingValues.getOrDefault(COLUMN_CUSTOMER_SHARDING_KEY, Collections.emptyList()));

        Set<Integer> suffixes = new HashSet<>();
        orderIds.forEach(id -> suffixes.add(extractSuffix(id)));
        customerIds.forEach(id -> suffixes.add(extractSuffix(id)));

        suffixes.forEach(suffix -> {
            int tableSuffix = suffix % availableTargetNames.size();
            availableTargetNames.stream()
                    .filter(table -> table.endsWith(String.valueOf(tableSuffix)))
                    .findFirst()
                    .ifPresent(result::add);
        });

        return result.isEmpty() ? availableTargetNames : result;
    }

    private Collection<String> handleRangeSharding(Collection<String> availableTargetNames, ComplexKeysShardingValue<Long> shardingValue) {

        Map<String, Range<Long>> rangeMap = shardingValue.getColumnNameAndRangeValuesMap();
        if (!rangeMap.containsKey(COLUMN_ORDER_SHARDING_KEY)) {
            return availableTargetNames;
        }

        Range<Long> range = rangeMap.get(COLUMN_ORDER_SHARDING_KEY);
        String lower = range.hasLowerBound() ? range.lowerEndpoint().toString() : "00";
        String upper = range.hasUpperBound() ? range.upperEndpoint().toString() : "99";

        int lowerSuffix = extractSuffix(lower);
        int upperSuffix = extractSuffix(upper);

        Set<Integer> tableSuffixes = new HashSet<>();
        int tableCount = availableTargetNames.size();

        // 处理跨100的情况（如95到05）
        if (lowerSuffix <= upperSuffix) {
            for (int i = lowerSuffix; i <= upperSuffix; i++) {
                tableSuffixes.add(i % tableCount);
            }
        } else {
            for (int i = lowerSuffix; i < 100; i++) {
                tableSuffixes.add(i % tableCount);
            }
            for (int i = 0; i <= upperSuffix; i++) {
                tableSuffixes.add(i % tableCount);
            }
        }

        return availableTargetNames.stream()
                .filter(table -> {
                    String suffixStr = table.substring(table.length() - 2);
                    int suffix = Integer.parseInt(suffixStr);
                    return tableSuffixes.contains(suffix);
                })
                .collect(Collectors.toList());
    }

    private List<String> convertToStringList(Collection<Long> values) {
        return values.stream()
                .map(Object::toString)
                .collect(Collectors.toList());
    }

    private int extractSuffix(String id) {
        if (id.length() >= 2) {
            return Integer.parseInt(id.substring(id.length() - 2));
        }
        // 处理不足两位的情况（前补零）
        return Integer.parseInt(String.format("%02d", Integer.parseInt(id)));
    }

    @Override
    public Properties getProps() {
        return null;
    }

    @Override
    public void init(Properties properties) {

    }
}